"""
假如我们有一个问题: 给出一段文本，使用一些关键词对它进行描述!
key是预先给出的提示性关键词， 整个的文本信息就相当于是query，value是你看到文本信息后得出的答案，
未经过学习前，key是我们只能看到的答案，因此key与value基本是相同的，
经过深入学习，我们开始能对query进行关键信息提取，
这就是注意力作用的过程， 通过这个过程，
我们最终得到的答案value发生了变化，
根据提示key生成了query的关键词表示方法，也就是另外一种特征表示方法.

一般形式：K = V ≠ Q
自注意力机制：K = V = Q
使用一般注意力机制，是使用不同于给定文本的关键词表示它. 而自注意力机制,
需要用给定文本自身来表达自己，也就是说你需要从给定文本中抽取关键词来表述它, 相当于对文本自身的一次特征提取.

"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.autograd import Variable

"""采用转置矩阵相乘的方法构造注意力机制函数"""
def attention(query, key, value, mask=None, dropout=None):
    """注意力机制的实现, 输入分别是query, key, value, mask: 掩码张量,
       dropout是nn.Dropout层的实例化对象, 默认为None"""
    # 在函数中, 首先取query的最后一维的大小, 一般情况下就等同于我们的词嵌入维度, 命名为d_k
    d_k = query.size(-1)
    # 按照注意力公式, 将query与key的转置相乘, 这里面key是将最后两个维度进行转置, 再除以缩放系数根号下d_k, 这种计算方法也称为缩放点积注意力计算.
    # 得到注意力得分张量scores
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # 接着判断是否使用掩码张量
    if mask is not None:
        # 使用tensor的masked_fill方法, 将掩码张量和scores张量每个位置一一比较, 如果掩码张量处为0
        # 则对应的scores张量用-1e9这个值来替换, 如下演示
        scores = scores.masked_fill(mask == 0, -1e9)

    # 对scores的最后一维进行softmax操作, 使用F.softmax方法, 第一个参数是softmax对象, 第二个是目标维度.
    # 这样获得最终的注意力张量
    p_attn = F.softmax(scores, dim = -1)

    # 之后判断是否使用dropout进行随机置0
    if dropout is not None:
        # 将p_attn传入dropout对象中进行'丢弃'处理
        p_attn = dropout(p_attn)

    # 最后, 根据公式将p_attn与value张量相乘获得最终的query注意力表示, 同时返回注意力张量
    return torch.matmul(p_attn, value), p_attn



if __name__ == '__main__':
    input1 = Variable(torch.randn(5, 5))
    mask = Variable(torch.zeros(5, 5))
    input1.masked_fill(mask == 0, -1e9)

    # query = key = value = pe_result
    # attn, p_attn = attention(query, key, value)





